import yaml
from langchain_core.documents import Document
from langchain_milvus import Milvus
from langchain_ollama import OllamaEmbeddings
from langchain_text_splitters import MarkdownTextSplitter

# 假设你的 Markdown 文件路径
md_file_path = r"C:\Desktop\各类知识\知识手册\2019-03-08_中科宇图科技股份有限公司精准治霾基础知识学习手册_1.0.0.md"

# 读取 Markdown 文件
with open(md_file_path, "r", encoding="utf-8") as f:
    md_content = f.read()

# 提取 YAML Front Matter
def extract_yaml_front_matter(md_text):
    """提取 YAML Front Matter 并返回元数据和正文内容"""
    if md_text.startswith("---"):
        parts = md_text.split("---", 2)  # 分割 YAML 部分
        if len(parts) > 2:
            yaml_content, markdown_body = parts[1], parts[2]
            metadata = yaml.safe_load(yaml_content)  # 解析 YAML
            return metadata, markdown_body
    return {}, md_text  # 如果没有 YAML Front Matter，则返回空字典

# 获取元数据和正文
metadata, md_body = extract_yaml_front_matter(md_content)

# 使用 MarkdownTextSplitter 进行分块
text_splitter = MarkdownTextSplitter(chunk_size=500, chunk_overlap=50)
chunks = text_splitter.split_text(md_body)

# 你的 embedding 模型（可以替换为自己的）
embedder = OllamaEmbeddings(model="nomic-embed-text:latest", base_url="http://192.168.7.3:11434")

# 连接 Milvus
vectorstore = Milvus(
    embedding_function=embedder,
    collection_name="trust",
    connection_args={"uri": "http://192.168.6.20:19530/ipp_air_general"},
    auto_id=True
)

# 处理并存储分块数据
docs = []
for chunk in chunks:
    doc_metadata = metadata.copy()  # 复制元数据
    doc_metadata["chunk_length"] = len(chunk)  # 添加 chunk 长度信息
    docs.append(Document(page_content=chunk, metadata=doc_metadata))  # 这里使用 Document 结构封装

# 将数据存入 Milvus
vectorstore.add_documents(docs)

print(f"成功存入 {len(docs)} 个文本块到 Milvus！")
